#include "ONNXClassifier.h"
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/dnn/dnn.hpp>
#include <vector>
#include <fstream>
#include <string>
#include <assert.h>
using namespace std;
using namespace cv;

/**
 * @brief Construct a new ONNXClassifier::ONNXClassifier object
 * 
 * @param model_path onnx模型文件路径
 * @param label_path label标签文本文件路径
 * @param _input_size 模型输入图像的尺寸
 */
ONNXClassifier::ONNXClassifier(const string &model_path, const string &label_path, Size _input_size) : default_mean(0.485, 0.456, 0.406),
                                                                                                       default_std(0.229, 0.224, 0.225), input_size(_input_size)
{
    if (!read_labels(label_path)) //调用读取label文本文件方法
    {                             //根据read_labels返回的布尔数抛出异常
        throw runtime_error("label read fail! Please correct your text file according to the standard mentioned in README.md");
    }
    net = dnn::readNet(model_path);                    //构造网络模型
    net.setPreferableBackend(dnn::DNN_BACKEND_OPENCV); //设置计算后台
    net.setPreferableTarget(dnn::DNN_TARGET_CPU);      //设置目标设备
}

/**
 * @brief 读取label文本文件--将其中的class_name存入labels成员数组变量，返回读取情况
 * 
 * @param label_path label文本文件的路径
 * @return true 读取label文本文件成功
 * @return false 读取label文本文件失败（为空或格式不规范）
 */
bool ONNXClassifier::read_labels(const string &label_path)
{
    ifstream ifs(label_path); //读取文本文件
    assert(ifs.is_open());    //根据文件打开情况来终止程序
    string line;              //存储文本文件的每行字符
    while (getline(ifs, line))
    {
        size_t index = line.find_first_of(':');   //根据每行字符的':'字符得到冒号左边的索引
        labels.push_back(line.substr(index + 1)); //复制冒号右边的字符串并增加到labels成员数组变量中
    }
    if (labels.size() > 0) //判断labels数组是否为空
        return true;       //返回true，即读取label并输入到labels数组成功
    else
        return false; //返回false，读取读取label并输入到labels数组失败（label文本文件不规范等）
}

/**
 * @brief 输入图像预处理（归一化、减法运算、除法运算等）
 * 
 * @param image 输入图像矩阵
 */
void ONNXClassifier::preprocess_input(Mat &image)
{
    image.convertTo(image, CV_32F, 1.0 / 255.0); //alpha为可选的比例因子，使用1.0/255.0的比例因子可将元素都归一化为0~1的数
    subtract(image, default_mean, image);        //图像减法运算
    divide(image, default_std, image);           //图像除法运算
}

/**
 * @brief 输入图像，调用预处理函数，调整图像尺寸，输入神经网络模型进行推理，得到置信度最高的索引后返回对应的labels标签
 * 
 * @param input_image 输入图像的矩阵
 * @param out_name 推理得到对应的label标签名
 */
void ONNXClassifier::Classify(const Mat &input_image, string &out_name)
{
    out_name.clear();
    Mat image = input_image.clone(); //将input_image赋值给开辟新内存空间的image
    preprocess_input(image);         //调用预处理图像方法
    //blobFromImage：将图像处理成网络模型的图像输入尺寸
    Mat input_blob = dnn::blobFromImage(image, 1.0, input_size, Scalar(0, 0, 0), true);
    net.setInput(input_blob);                                             //设置网络模型的输入
    const vector<String> &out_names = net.getUnconnectedOutLayersNames(); //得到网络模型的每一层
    Mat out_tensor = net.forward(out_names[0]);                           //向前传播，得到结果张量out_tensor
    Point maxLoc;                                                         //存储置信度最高的索引
    minMaxLoc(out_tensor, (double *)0, (double *)0, (Point *)0, &maxLoc); //置信度最高对应的索引
    out_name = labels[maxLoc.x];                                          //根据置信度最高的索引和label标签返回推理结果
}